#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from struct import pack, unpack

from ambari_ws4py.exc import FrameTooLargeException, ProtocolException
from ambari_ws4py.compat import py3k, ord, range

# Frame opcodes defined in the spec.
OPCODE_CONTINUATION = 0x0
OPCODE_TEXT = 0x1
OPCODE_BINARY = 0x2
OPCODE_CLOSE = 0x8
OPCODE_PING = 0x9
OPCODE_PONG = 0xA

__all__ = ["Frame"]


class Frame(object):
  def __init__(
    self, opcode=None, body=b"", masking_key=None, fin=0, rsv1=0, rsv2=0, rsv3=0
  ):
    """
    Implements the framing protocol as defined by RFC 6455.

    .. code-block:: python
       :linenos:

       >>> test_mask = 'XXXXXX' # perhaps from os.urandom(4)
       >>> f = Frame(OPCODE_TEXT, 'hello world', masking_key=test_mask, fin=1)
       >>> bytes = f.build()
       >>> bytes.encode('hex')
       '818bbe04e66ad6618a06d1249105cc6882'
       >>> f = Frame()
       >>> f.parser.send(bytes[0])
       1
       >>> f.parser.send(bytes[1])
       4

    .. seealso:: Data Framing http://tools.ietf.org/html/rfc6455#section-5.2
    """
    if not isinstance(body, bytes):
      raise TypeError("The body must be properly encoded")

    self.opcode = opcode
    self.body = body
    self.masking_key = masking_key
    self.fin = fin
    self.rsv1 = rsv1
    self.rsv2 = rsv2
    self.rsv3 = rsv3
    self.payload_length = len(body)

    self._parser = None

  @property
  def parser(self):
    if self._parser is None:
      self._parser = self._parsing()
      # Python generators must be initialized once.
      next(self.parser)
    return self._parser

  def _cleanup(self):
    if self._parser:
      self._parser.close()
      self._parser = None

  def build(self):
    """
    Builds a frame from the instance's attributes and returns
    its bytes representation.
    """
    header = b""

    if self.fin > 0x1:
      raise ValueError("FIN bit parameter must be 0 or 1")

    if 0x3 <= self.opcode <= 0x7 or 0xB <= self.opcode:
      raise ValueError("Opcode cannot be a reserved opcode")

    ## +-+-+-+-+-------+
    ## |F|R|R|R| opcode|
    ## |I|S|S|S|  (4)  |
    ## |N|V|V|V|       |
    ## | |1|2|3|       |
    ## +-+-+-+-+-------+
    header = pack(
      "!B",
      (
        (self.fin << 7)
        | (self.rsv1 << 6)
        | (self.rsv2 << 5)
        | (self.rsv3 << 4)
        | self.opcode
      ),
    )

    ##                 +-+-------------+-------------------------------+
    ##                 |M| Payload len |    Extended payload length    |
    ##                 |A|     (7)     |             (16/63)           |
    ##                 |S|             |   (if payload len==126/127)   |
    ##                 |K|             |                               |
    ## +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
    ## |     Extended payload length continued, if payload len == 127  |
    ## + - - - - - - - - - - - - - - - +-------------------------------+
    if self.masking_key:
      mask_bit = 1 << 7
    else:
      mask_bit = 0

    length = self.payload_length
    if length < 126:
      header += pack("!B", (mask_bit | length))
    elif length < (1 << 16):
      header += pack("!B", (mask_bit | 126)) + pack("!H", length)
    elif length < (1 << 63):
      header += pack("!B", (mask_bit | 127)) + pack("!Q", length)
    else:
      raise FrameTooLargeException()

    ## + - - - - - - - - - - - - - - - +-------------------------------+
    ## |                               |Masking-key, if MASK set to 1  |
    ## +-------------------------------+-------------------------------+
    ## | Masking-key (continued)       |          Payload Data         |
    ## +-------------------------------- - - - - - - - - - - - - - - - +
    ## :                     Payload Data continued ...                :
    ## + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
    ## |                     Payload Data continued ...                |
    ## +---------------------------------------------------------------+
    body = self.body
    if not self.masking_key:
      return bytes(header + body)

    return bytes(header + self.masking_key + self.mask(body))

  def _parsing(self):
    """
    Generator to parse bytes into a frame. Yields until
    enough bytes have been read or an error is met.
    """
    buf = b""
    some_bytes = b""

    # yield until we get the first header's byte
    while not some_bytes:
      some_bytes = yield 1

    first_byte = (
      some_bytes[0] if isinstance(some_bytes, bytearray) else ord(some_bytes[0])
    )
    # frame-fin = %x0 ; more frames of this message follow
    #           / %x1 ; final frame of this message
    self.fin = (first_byte >> 7) & 1
    self.rsv1 = (first_byte >> 6) & 1
    self.rsv2 = (first_byte >> 5) & 1
    self.rsv3 = (first_byte >> 4) & 1
    self.opcode = first_byte & 0xF

    # frame-rsv1 = %x0 ; 1 bit, MUST be 0 unless negotiated otherwise
    # frame-rsv2 = %x0 ; 1 bit, MUST be 0 unless negotiated otherwise
    # frame-rsv3 = %x0 ; 1 bit, MUST be 0 unless negotiated otherwise
    if self.rsv1 or self.rsv2 or self.rsv3:
      raise ProtocolException()

    # control frames between 3 and 7 as well as above 0xA are currently reserved
    if 2 < self.opcode < 8 or self.opcode > 0xA:
      raise ProtocolException()

    # control frames cannot be fragmented
    if self.opcode > 0x7 and self.fin == 0:
      raise ProtocolException()

    # do we already have enough some_bytes to continue?
    some_bytes = some_bytes[1:] if some_bytes and len(some_bytes) > 1 else b""

    # Yield until we get the second header's byte
    while not some_bytes:
      some_bytes = yield 1

    second_byte = (
      some_bytes[0] if isinstance(some_bytes, bytearray) else ord(some_bytes[0])
    )
    mask = (second_byte >> 7) & 1
    self.payload_length = second_byte & 0x7F

    # All control frames MUST have a payload length of 125 some_bytes or less
    if self.opcode > 0x7 and self.payload_length > 125:
      raise FrameTooLargeException()

    if some_bytes and len(some_bytes) > 1:
      buf = some_bytes[1:]
      some_bytes = buf
    else:
      buf = b""
      some_bytes = b""

    if self.payload_length == 127:
      # This will compute the actual application data size
      if len(buf) < 8:
        nxt_buf_size = 8 - len(buf)
        some_bytes = yield nxt_buf_size
        some_bytes = buf + (some_bytes or b"")
        while len(some_bytes) < 8:
          b = yield 8 - len(some_bytes)
          if b is not None:
            some_bytes = some_bytes + b
        if len(some_bytes) > 8:
          buf = some_bytes[8:]
          some_bytes = some_bytes[:8]
      else:
        some_bytes = buf[:8]
        buf = buf[8:]
      extended_payload_length = some_bytes
      self.payload_length = unpack("!Q", extended_payload_length)[0]
      if self.payload_length > 0x7FFFFFFFFFFFFFFF:
        raise FrameTooLargeException()
    elif self.payload_length == 126:
      if len(buf) < 2:
        nxt_buf_size = 2 - len(buf)
        some_bytes = yield nxt_buf_size
        some_bytes = buf + (some_bytes or b"")
        while len(some_bytes) < 2:
          b = yield 2 - len(some_bytes)
          if b is not None:
            some_bytes = some_bytes + b
        if len(some_bytes) > 2:
          buf = some_bytes[2:]
          some_bytes = some_bytes[:2]
      else:
        some_bytes = buf[:2]
        buf = buf[2:]
      extended_payload_length = some_bytes
      self.payload_length = unpack("!H", extended_payload_length)[0]

    if mask:
      if len(buf) < 4:
        nxt_buf_size = 4 - len(buf)
        some_bytes = yield nxt_buf_size
        some_bytes = buf + (some_bytes or b"")
        while not some_bytes or len(some_bytes) < 4:
          b = yield 4 - len(some_bytes)
          if b is not None:
            some_bytes = some_bytes + b
        if len(some_bytes) > 4:
          buf = some_bytes[4:]
      else:
        some_bytes = buf[:4]
        buf = buf[4:]
      self.masking_key = some_bytes

    if len(buf) < self.payload_length:
      nxt_buf_size = self.payload_length - len(buf)
      some_bytes = yield nxt_buf_size
      some_bytes = buf + (some_bytes or b"")
      while len(some_bytes) < self.payload_length:
        l = self.payload_length - len(some_bytes)
        b = yield l
        if b is not None:
          some_bytes = some_bytes + b
    else:
      if self.payload_length == len(buf):
        some_bytes = buf
      else:
        some_bytes = buf[: self.payload_length]

    self.body = some_bytes
    yield

  def mask(self, data):
    """
    Performs the masking or unmasking operation on data
    using the simple masking algorithm:

    ..
       j                   = i MOD 4
       transformed-octet-i = original-octet-i XOR masking-key-octet-j

    """
    masked = bytearray(data)
    if py3k:
      key = self.masking_key
    else:
      key = map(ord, self.masking_key)
    for i in range(len(data)):
      masked[i] = masked[i] ^ key[i % 4]
    return masked

  unmask = mask
